# -*- coding:utf-8 -*-
import base64

from PIL import Image
from io import BytesIO
import pytesseract


from spider.common.util.JSONUtil import JsonUtil


class OcrUtil:
    def __init__(self):
        self.jsonUtil=JsonUtil()

    def image_to_base64(self, image_path):
        with open(image_path, "rb") as img_file:
            base64_data = base64.b64encode(img_file.read())
            return base64_data.decode('utf-8')

    def base64_to_image(self, base64_data, image_path='demo.png'):
        image_data = base64.b64decode(base64_data)
        image = Image.open(BytesIO(image_data))
        image.save(image_path)
        return image_path

    def recognize_text(self, base64_str, lang='eng'):
        """
        识别Base64图片中的文字

        :param base64_str: Base64编码的图片字符串
        :param lang: 识别语言，默认为英文('eng')，中文可用'chi_sim'
        :return: 识别出的文字
        """
        try:
            # 将Base64转换为图片
            image = self.base64_to_image(base64_str)

            # 进行OCR识别
            text = pytesseract.image_to_string(image, lang=lang)

            # 去除多余的空白字符
            return text.strip()
        except Exception as e:
            raise RuntimeError(f"文字识别失败: {str(e)}")


if __name__ == "__main__":
    ocr_client = OcrUtil()
    data=ocr_client.image_to_base64('./a.png')
    code=ocr_client.recognize_text(base64_str=data)
    print(code)